# -*- UTF-8 -*-
import sys
import xml.etree.ElementTree as ET
import pandas as pd
import numpy as np


def parse_basic(xml_file, attr_list, csv_path):
    parser = ET.XMLParser(encoding="utf-8")
    tree = ET.parse(xml_file, parser=parser)
    print("Parse xml succeed!!!")

    rows = tree.findall('row')

    output_num = len(rows)
    # 预先分配大小
    print("Creating DataFrame:")
    arr = np.empty(output_num, dtype=map)
    for i in range(output_num):
        arr[i] = rows[i].attrib

    df = pd.DataFrame(arr.tolist(), columns=attr_list)

    df.to_csv(csv_path, index=False)
    print("Written to csv!")


def parse_posts(questions_attrs, answers_attrs):
    parser = ET.XMLParser(encoding="utf-8")
    tree = ET.parse("/home/LAB/caipb/data/Posts.xml", parser=parser)
    print("Parse xml succeed!!!")

    rows = tree.findall('row')

    posts_num = 41782536
    answer_num = 26000000
    question_num = 17000000

    # 预先分配大小
    print("Creating DataFrame:")
    answer_arr = np.empty(answer_num, dtype=map)
    question_arr = np.empty(question_num, dtype=map)
    q_index = 0
    a_index = 0

    # 分类
    for i in range(posts_num):
        if rows[i].attrib['PostTypeId'] == '1':
            question_arr[q_index] = rows[i].attrib
            q_index = q_index + 1
        elif rows[i].attrib['PostTypeId'] == '2':
            answer_arr[a_index] = rows[i].attrib
            a_index = a_index + 1

    print("Answers count: " + a_index)
    print("Questions count: ", q_index)

    question_df = pd.DataFrame(question_arr.tolist(), columns=questions_attrs)
    answer_df = pd.DataFrame(answer_arr.tolist(), columns=answers_attrs)

    question_df.to_csv("/home/incoming/caipb/Questions.csv", index=False)
    answer_df.to_csv("/home/incoming/caipb/Answers.csv", index=False)

    print("Written to csv!")


if __name__ == '__main__':
    answer_attrs_list = ['ParentId',
                         'CreationDate',
                         'CommunityOwnedDate',
                         'ViewCount',
                         'Title',
                         'LastEditorUserId',
                         'OwnerDisplayName',
                         'LastEditorDisplayName',
                         'Score',
                         'OwnerUserId',
                         'LastActivityDate',
                         'FavoriteCount',
                         'CommentCount',
                         'Id',
                         'LastEditDate',
                         'ClosedDate',
                         'Body']
    question_attrs_list = ['CreationDate',
                           'CommunityOwnedDate',
                           'ViewCount',
                           'Title',
                           'LastEditorUserId',
                           'OwnerDisplayName',
                           'LastEditorDisplayName',
                           'Score',
                           'OwnerUserId',
                           'LastActivityDate',
                           'FavoriteCount',
                           'AnswerCount',
                           'CommentCount',
                           'Id',
                           'LastEditDate',
                           'ClosedDate',
                           'AcceptedAnswerId',
                           'Body',
                           'Tags']
    user_list = ['CreationDate',
                 'Views',
                 'AccountId',
                 'AboutMe',
                 'LastAccessDate',
                 'UpVotes',
                 'ProfileImageUrl',
                 'Card',
                 'DisplayName',
                 'Id',
                 'Reputation',
                 'DownVotes',
                 'Age',
                 'Location'
                 ]

    tag_attrs_list = ['Id', 'TagName', 'Count', 'ExcerptPostId', 'WikiPostId']
    badge_attrs_list = ['UserId', 'Class', 'Id', 'TagBased', 'Date', 'Name']

    comment_attrs_list = ['CreationDate', 'Score', 'UserId', 'Text', 'Id', 'PostId']

    # user_num =
    # parse_basic("E:\Folder Storing Codes\PycharmProjects\Titanic\Tags.xml", tag_attrs_list)
    parse_basic("/home/LAB/caipb/data/Tags.xml", tag_attrs_list, "/home/incoming/caipb/Tags.csv")
